/*
 * Copyright (c) 2018. the original author or authors.
 * Kunlun is licensed under the "LICENSE" file in the project's root directory.
 */

package kunlun.ai.support.model;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;

/**
 * The AI chat response.
 * @author Kahle
 */
public class ChatResponse implements Serializable {
    /**
     * A unique identifier for the chat.
     */
    private String id;
    /**
     * The model used for the chat.
     */
    private String model;
    /**
     * A list of chat choices.
     * Can be more than one if n is greater than 1.
     */
    private List<Choice> choices;
    /**
     * The usage statistics for the completion request.
     */
    private Usage usage;

    public ChatResponse(String model) {

        this.model = model;
    }

    public ChatResponse() {

    }

    public String getId() {

        return id;
    }

    public void setId(String id) {

        this.id = id;
    }

    public String getModel() {

        return model;
    }

    public void setModel(String model) {

        this.model = model;
    }

    public List<Choice> getChoices() {

        return choices;
    }

    public void setChoices(List<Choice> choices) {

        this.choices = choices;
    }

    public Usage getUsage() {

        return usage;
    }

    public void setUsage(Usage usage) {

        this.usage = usage;
    }

    /**
     * The AI chat response choice.
     * @author Kahle
     */
    public static class Choice implements Serializable {
        /**
         * The reason the model stopped generating tokens.
         */
        private String  finishReason;
        /**
         * A chat message generated by the model.
         */
        private Message message;
        /**
         * The index of the choice in the list of choices.
         */
        private Integer index;

        public Choice(Integer index, Message message, String finishReason) {
            this.finishReason = finishReason;
            this.message = message;
            this.index = index;
        }

        public Choice(Message message, String finishReason) {
            this.finishReason = finishReason;
            this.message = message;
        }

        public Choice() {

        }

        public Integer getIndex() {

            return index;
        }

        public void setIndex(Integer index) {

            this.index = index;
        }

        public Message getMessage() {

            return message;
        }

        public void setMessage(Message message) {

            this.message = message;
        }

        public String getFinishReason() {

            return finishReason;
        }

        public void setFinishReason(String finishReason) {

            this.finishReason = finishReason;
        }
    }

    /**
     * The AI chat response builder.
     * @author Kahle
     */
    public static class Builder implements kunlun.core.Builder {

        public static Builder of(String model) {

            return of().setModel(model);
        }

        public static Builder of() {

            return new Builder();
        }

        private String id;
        private String model;
        private final List<Choice> choices = new ArrayList<Choice>();
        private Usage usage;

        public String getId() {

            return id;
        }

        public Builder setId(String id) {
            this.id = id;
            return this;
        }

        public String getModel() {

            return model;
        }

        public Builder setModel(String model) {
            this.model = model;
            return this;
        }

        public List<Choice> getChoices() {

            return choices;
        }

        public Builder addChoices(List<Choice> choices) {
            this.choices.addAll(choices);
            return this;
        }

        public Builder addChoice(Choice choice) {
            this.choices.add(choice);
            return this;
        }

        public Builder addChoice(Integer index, Message message, String finishReason) {
            this.choices.add(new Choice(index, message, finishReason));
            return this;
        }

        public Builder addChoice(Message message, String finishReason) {
            this.choices.add(new Choice(message, finishReason));
            return this;
        }

        public Usage getUsage() {

            return usage;
        }

        public Builder setUsage(Usage usage) {
            this.usage = usage;
            return this;
        }

        @Override
        public ChatResponse build() {
            ChatResponse response = new ChatResponse(model);
            response.setId(id);
            response.setChoices(choices);
            response.setUsage(usage);
            return response;
        }
    }

}
